2157. Sum

 

Roman’s parents gave him an undirected connected weighted graph with n vertices and n – 1 edges. Roman wants to know the sum of all the paths’ lengths in this graph. Let’s consider the path’s length as the sum of all the edges that lie on this path. Roman said that the path from u to v is the same as from v to u, so he doesn’t distinguish them.

 

Input. The first line contains the single integer number n (2 ≤ n ≤ 105) – the number of vertices in the graph. It is followed by n – 1 lines containing the description of the edges. Every line of description consists of three integers: the vertices connected by the edge (labeled from 1 to n) and the edge’s weight.

 

Output. Print the sum of all the paths’ lengths in the given graph modulo 109.

 

Sample input 1

Sample output 1

3

1 2 1

1 3 3

8

 

 

Sample input 2

Sample output 2

6

1 2 5

1 3 1

2 4 2

2 5 4

2 6 3

90

 

In the first sample all paths are: 1 – 2, 1 – 3, 2 – 1 – 3 and their lengths’ sum is 1 + 3 + 4 = 8.

 

 

SOLUTION

graphs, depth first search, tree

 

Algorithm analysis

Start the depth-first search from some vertex of the tree. Consider an edge (u, v) of a tree with weight w. Let the number of vertices in the subtree with the root v be equal to c. Then the tree contains c vertices on one side of the edge, and n – c vertices on the other. The edge (u, v) appears in c * (n – c) different paths. Therefore, the contribution of this edge to the sum of the lengths of all paths is c * (n – c) * w. It remains to iterate through all the edges with a depth first search and sum their contributions to the total length.

 

 

Example

The graphs given in the samples have the following form:

 

Consider the contributions of the edges to the total sum of lengths of all paths in the first example.

The edge (1, 2) of weight 1 belongs to two paths:

·        1 – 2;

·        2 – 1 – 3;

Its contribution to the total sum is 1 * 2 * 1 = 2.

The edge (1, 3) of weight 3 belongs to two paths:

·        1 – 3;

·        2 – 1 – 3;

Its contribution to the total sum is 2 * 1 * 3 = 6.

The sum of the lengths of all paths is 2 + 6 = 8.

 

In the second sample, consider the contribution of the edge (1, 2) of weight 5 to the total sum of the lengths of all paths.

The number of vertices in the subtree with root 2 is c = 4 (including the vertex 2 itself). Then on the other side of the edge (1, 2) there are n – c = 6 – 4 = 2 vertices. Therefore, the number of different paths containing the edge (1, 2) is equal to c * (n – c) = 4 * 2 = 8. Indeed, such paths are

1 – 2, 1 – 2 – 4, 1 – 2 – 5, 1 – 2 – 6,

3 – 1 – 2, 3 – 1 – 2 – 4, 3 – 1 – 2 – 5, 3 – 1 – 2 – 6

The contribution of the edge (1, 2) of weight 5 to the sum of the lengths of all paths is c * (n – c) * w = 4 * 2 * 5 = 40.

 

Algorithm realization

Store the input weighted graph in the adjacency list g.

 

#define MOD 1000000000;

vector<vector<pair<int,int> > > g;

 

Function dfs implements a depth-first search from vertex v and returns the number of vertices in the subtree rooted in v (including v itself). These vertices are counted in the variable cnt. Initially set cnt = 1, assuming that the vertex v itself is included in the subtree.

 

int dfs(int v, int p = -1)

{

  int cnt = 1;

  for (auto edge : g[v])

  {

    int to = edge.first;

    int w = edge.second;

 

Consider an edge (v, to) with weight w. Calculate the number of vertices c in the subtree rooted at to. Thus, the tree contains c vertices on one side of the edge, and n – c vertices on the other. The edge (v, to) is included in c * (n – c) different paths. Therefore, the contribution of this edge to the sum of the lengths of all paths is c * (n – c) * w.

 

    if (to != p)

    {

      int c = dfs(to, v);

      res = (res + 1LL * c * (n - c) * w) % MOD;

      cnt += c;

    }

  }

  return cnt;

}

 

The main part of the program. Read the weighted graph into the adjacency list g.

 

scanf("%d", &n);

g.resize(n + 1);

for(i = 1; i < n; i++)

{

  scanf("%d %d %d", &u, &v, &d);

  g[u].push_back({v, d});

  g[v].push_back({u, d});

}

 

Run a depth-first search from vertex 1. The vertices of the graph are numbered from 1 to n.

 

dfs(1);

 

Print the answer.

 

printf("%lld\n",res);

 

Java realization

 

import java.util.*;

import java.io.*;

 

class Edge

{

  int to;

  int weight;

  public Edge(int to, int weight)

  {

    this.to = to;

    this.weight = weight;

  }

}

 

public class Main

{

  static int MOD = 1000000000;

  static ArrayList<Edge>[] g;    

  static int n, m;

  static long res;

 

  static int dfs(int v, int p)

  {

    int cnt = 1;

    for(Edge edge : g[v])

    {

      int to = edge.to;

      int w = edge.weight;

      if (to != p)

      {

        int c = dfs(to, v);

        res = (res + 1L * c * (n - c) * w) % MOD;

        cnt += c;

      }

    }

    return cnt;

  }

 

  public static void main(String[] args)

  {

    FastScanner con = new FastScanner(System.in);

    n = con.nextInt();

    g = new ArrayList[n+1]; // unchecked

 

    for (int i = 0; i <= n; i++)

      g[i] = new ArrayList<Edge>();

 

    for (int i = 1; i < n; i++)

    {

      int u = con.nextInt();

      int v = con.nextInt();

      int d = con.nextInt();

      g[u].add(new Edge(v,d));

      g[v].add(new Edge(u,d));

    }

   

    dfs(1,-1);

    System.out.println(res);   

  }

}  

 

class FastScanner

{

  BufferedReader br;

  StringTokenizer st;

 

  public FastScanner(InputStream inputStream)

  {

    br = new BufferedReader(new InputStreamReader(inputStream));

    st = new StringTokenizer("");

  }

 

  public String next()

  {

    while (!st.hasMoreTokens())

    {

      try

      {

        st = new StringTokenizer(br.readLine());

      } catch (Exception e)

      {

        return null;

      }

    }

    return st.nextToken();

  }

 

  public int nextInt()

  {

    return Integer.parseInt(next());

  }

}

 

Python realization

 

MOD = 1000000000

 

def dfs(v, p = -1):

  global res

  cnt = 1

  for to, w in g[v]:

    if to != p:

      c = dfs(to, v)

      res = (res + c * (n - c) * w) % MOD

      cnt += c

  return cnt

 

n = int(input())

g = [[] for _ in range(n + 1)]

res = 0

 

for _ in range(n - 1):

  u, v, d = map(int, input().split())

  g[u].append((v, d))

  g[v].append((u, d))

 

dfs(1)

 

print(res)